import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from models.fcn8s import FCN8s
from utils.datasets import VOCSegmentation
from torchvision import transforms
import random
import matplotlib.pyplot as plt
import numpy as np

# 数据预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 数据集和数据加载器
root = 'data/VOCdevkit/VOC2012'
image_set_file = os.path.join(root, 'ImageSets', 'Segmentation', 'trainval.txt')
with open(image_set_file) as f:
    all_image_ids = [line.strip() for line in f]

random.shuffle(all_image_ids)
train_size = int(0.8 * len(all_image_ids))
train_image_ids = all_image_ids[:train_size]
test_image_ids = all_image_ids[train_size:]

train_dataset = VOCSegmentation(root, image_set='trainval', transform=transform, image_ids=train_image_ids)
test_dataset = VOCSegmentation(root, image_set='trainval', transform=transform, image_ids=test_image_ids)

train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=4, shuffle=False)

# 模型、损失函数和优化器
model = FCN8s(n_class=21)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练循环
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
num_epochs = 10

#标签打印测试begin
# 检查部分标签值
for i in range(10):
    image, label = train_dataset[i]
    print(f"Sample {i} label min: {label.min()}, max: {label.max()}")
#标签打印测试end

for epoch in range(num_epochs):
    running_loss = 0.0
    for images, labels in train_loader:
        images = images.to(device)
        labels = labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        # print("trainoutput",outputs.size())

        # 提取 labels 的空间维度
        if len(labels.shape) == 4:  # 如果 labels 形状是 (N, 1, H, W)
            output_size = labels.shape[2:]
        elif len(labels.shape) == 3:  # 如果 labels 形状是 (N, H, W)
            output_size = labels.shape[1:]
        else:
            raise ValueError("labels 形状不符合预期，请检查数据加载部分。")

        # 进行插值操作
        outputs = torch.nn.functional.interpolate(outputs, size=output_size, mode='bilinear', align_corners=True)

        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        print("trainnig...")
    print(f'Epoch {epoch + 1}/{num_epochs}, Loss: {running_loss / len(train_loader)}')

print('Training finished.')

# 保存训练好的模型
torch.save(model.state_dict(), 'fcn8s_model.pth')

# 测试循环
model.eval()  # 设置模型为评估模式
total_loss = 0.0
with torch.no_grad():  # 不需要计算梯度
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)

        outputs = model(images)
        # 上采样 outputs 到和 labels 相同的尺寸
        outputs = torch.nn.functional.interpolate(outputs, size=labels.shape[1:], mode='bilinear', align_corners=True)

        loss = criterion(outputs, labels)
        total_loss += loss.item()

    average_loss = total_loss / len(test_loader)
    print(f'Test Loss: {average_loss}')

# 可视化单张分割效果图
image, label = test_dataset[0]
image = image.unsqueeze(0).to(device)  # 添加一个维度，模拟批量输入并移到设备上

# 进行预测
with torch.no_grad():
    output = model(image)

    # 上采样 output 到和 label 相同的尺寸(使用双线性插值上采样)
    output = torch.nn.functional.interpolate(output, size=label.shape, mode='bilinear', align_corners=True)

# 获取预测结果
prediction = output.argmax(dim=1).squeeze().cpu().numpy()

# 可视化显示
def show_segmentation(image, label, prediction):
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    # 显示原始图像
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    image = image.squeeze(0).cpu().numpy()
    image = (image * std[:, None, None]) + mean[:, None, None]
    image = np.transpose(image, (1, 2, 0))
    axes[0].imshow(np.clip(image, 0, 1))
    axes[0].set_title('Original Image')
    axes[0].axis('off')

    # 显示真实标签
    axes[1].imshow(label.cpu().numpy())
    axes[1].set_title('Ground Truth')
    axes[1].axis('off')

    # 显示预测结果
    axes[2].imshow(prediction)
    axes[2].set_title('Prediction')
    axes[2].axis('off')

    plt.show()

# 显示分割结果
show_segmentation(image, label, prediction)
#gitceshi 